import os
from typing import Optional

import pandas as pd
from catboost import CatBoostClassifier, Pool
from selector.selector import XFileSelector
from constant import model_path
from src.train import FlowPipline, FlowInterrupt, TrainParam, FlowEnum
from src.util.common_util import predict_sample, test_evaluate, get_exclude_cols


@FlowPipline.bind(FlowEnum.step1_read_data, is_train=True, is_predict=True)
def read_data(df: Optional[pd.DataFrame | None], params: TrainParam):
    selector = XFileSelector(select_file=True, is_multi_file=True)
    files = selector.select_files()
    df = pd.DataFrame()
    if files:
        for path in files:
            df = pd.concat([df, pd.read_csv(path)])
    df.reset_index(inplace=True)
    if len(df) == 0:
        raise FlowInterrupt("读取数据为空")
    params.console.print(f"导入:{len(df)}条数据")
    return df


@FlowPipline.bind(FlowEnum.step10_train_or_predict, is_train=True, is_predict=False)
def train(df, params):
    # df.to_csv("prehandle.csv", index=False)
    df_cols = []
    for col in df.columns.to_list():
        if col not in get_exclude_cols():
            df_cols.append(col)
    df = df[df_cols]
    x_resampled, y_resampled, x_test, y_test, test_size, sample_type = predict_sample(df)
    model = CatBoostClassifier(iterations=1000, depth=14, learning_rate=0.1,
                               loss_function="Logloss", use_best_model=True,
                               eval_metric="F1")
    model.fit(Pool(x_resampled, label=y_resampled), eval_set=[(x_test, y_test)])
    accuracy, precision, recall, f1 = test_evaluate(model, x_test, y_test)
    os.makedirs(os.path.dirname(model_path), exist_ok=True)
    model.save_model(model_path)
    return None


@FlowPipline.bind(FlowEnum.step10_train_or_predict, is_train=False, is_predict=True)
def predict(df, params):
    df_cols = []
    for col in df.columns.to_list():
        if col not in get_exclude_cols():
            df_cols.append(col)
    df = df[df_cols]
    model = CatBoostClassifier()
    model.load_model(model_path)
    x_data = df[[item for item in df.columns.tolist() if item != 'is_dangerous']]
    y_pred = model.predict(x_data)
    result_normal = sum(1 for item in y_pred if f'{item}' == '0')
    result_abnormal = sum(1 for item in y_pred if f'{item}' == '1')
    params.console.print(f"[文件预测]文件预测结果：异常有{result_abnormal}个|正常有{result_normal}个")
    if result_abnormal <= 0 or result_normal <= 0:
        params.console.print("[文件预测]预测结果中没有正常或异常数据")
        return
    origin_df = params.origin_df

